"""
单张图片替换PDF指定位置：使用JSON中的bbox_page_pixel参数
"""

import os
import io
import json
from PIL import Image
import fitz  # PyMuPDF

# ========= 配置 =========
PDF_PATH = " "  # 替换为PDF文件路径
IMAGE_PATH = ""  # 替换为图片路径
JSON_PATH = " "  # 替换为JSON文件路径
OUTPUT_PATH = " "  # 设置为None使用默认命名，或指定输出路径

# 输入参数
PAPER_ID = "2010072785"  # paper_id
PANEL_ID = "2010072785_6_1_1"  # panel_id
PAGE_PX_PER_PT = 1.0  # 像素到点的转换比率

# ========= 基础工具 =========
def safe_overwrite_pdf(doc: fitz.Document, out_pdf: str) -> None:
    """安全覆盖保存PDF，避免repaired PDF报错"""
    try:
        if hasattr(doc, "can_save_incrementally") and doc.can_save_incrementally():
            doc.save(out_pdf, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
            return
        doc.save(out_pdf, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
        return
    except Exception:
        tmp = out_pdf + ".tmp"
        doc.save(tmp, garbage=4, deflate=True, encryption=fitz.PDF_ENCRYPT_KEEP)
        try:
            os.replace(tmp, out_pdf)
        except Exception:
            import shutil
            shutil.move(tmp, out_pdf)

def crop_and_resize_to_bytes(ori_img_path: str, out_w: int, out_h: int) -> bytes:
    """将图片裁剪并调整到指定尺寸，返回字节流"""
    with Image.open(ori_img_path) as img:
        # 保持宽高比的情况下居中裁剪
        target_aspect = out_w / float(out_h) if out_h != 0 else 1
        w, h = img.size
        
        if w / float(h) > target_aspect:
            new_w = int(round(h * target_aspect))
            x0 = max(0, (w - new_w) // 2)
            img = img.crop((x0, 0, x0 + new_w, h))
        else:
            new_h = int(round(w / target_aspect)) if target_aspect != 0 else h
            y0 = max(0, (h - new_h) // 2)
            img = img.crop((0, y0, w, y0 + new_h))
        
        # 调整到精确尺寸
        img = img.convert("RGB").resize((max(1, out_w), max(1, out_h)), Image.Resampling.LANCZOS)
        
        # 转换为字节流
        bio = io.BytesIO()
        img.save(bio, format="PNG", optimize=True)
        return bio.getvalue()

def page_pixel_to_pt(bbox_px: list) -> list:
    """将像素坐标转换为点坐标"""
    k = PAGE_PX_PER_PT
    return [bbox_px[0] / k, bbox_px[1] / k, bbox_px[2] / k, bbox_px[3] / k]

def find_bbox_in_json(json_data, paper_id, panel_id=None):
    """
    在JSON数据中查找指定论文和面板的bbox_page_pixel
    
    参数:
        json_data: JSON数据
        paper_id: 论文ID
        panel_id: 面板ID（可选，如果为None则使用第一个panel）
    
    返回:
        (bbox, page_index): 边界框和页面索引
    """
    # 首先找到指定的论文
    papers = json_data.get("papers", [])
    target_paper = None
    for paper in papers:
        if str(paper.get("paper_id")) == str(paper_id):
            target_paper = paper
            break
    
    if not target_paper:
        print(f"错误：未找到论文{paper_id}")
        return None, None
    
    print(f"找到论文{paper_id}")
    
    # 在指定论文中查找panel的bbox_page_pixel
    pages = target_paper.get("pages", [])
    for page in pages:
        page_index = page.get("page_index", 1)  # JSON中是1-based
        figures = page.get("figures", [])
        
        for fig in figures:
            panels = fig.get("panels", [])
            for panel in panels:
                if "bbox_page_pixel" in panel:
                    current_panel_id = panel.get("panel_id", "")
                    
                    # 如果指定了panel_id，检查是否匹配
                    if panel_id and current_panel_id != panel_id:
                        continue
                    
                    bbox = panel["bbox_page_pixel"]
                    print(f"找到panel坐标: 第{page_index}页, bbox={bbox}")
                    print(f"  panel_id: {current_panel_id}")
                    return bbox, page_index - 1  # 转换为0-based
    
    if panel_id:
        print(f"错误：在论文{paper_id}中未找到面板{panel_id}")
    else:
        print(f"错误：在论文{paper_id}中未找到任何panel")
    return None, None

# ========= 核心功能 =========
def replace_image_in_pdf(pdf_path: str, image_path: str, bbox_px: list, page_index: int, 
                         output_path: str = None):
    """
    将图片替换到PDF的指定位置
    
    参数:
        pdf_path: 原始PDF文件路径
        image_path: 要插入的图片路径
        bbox_px: 替换区域的边界框 [x0, y0, x1, y1] (像素单位)
        page_index: 页面索引(0-based)
        output_path: 输出PDF路径(默认为原PDF名加_replaced后缀)
    """
    # 设置输出路径
    if output_path is None:
        base, ext = os.path.splitext(pdf_path)
        output_path = f"{base}_replaced{ext}"
    
    # 打开PDF文档
    doc = fitz.open(pdf_path)
    
    # 检查页面索引是否有效
    if page_index < 0 or page_index >= len(doc):
        raise ValueError(f"页面索引 {page_index} 无效，文档只有 {len(doc)} 页")
    
    # 获取指定页面
    page = doc[page_index]
    
    # 将像素坐标转换为点坐标
    bbox_pt = page_pixel_to_pt(bbox_px)
    rect = fitz.Rect(bbox_pt)
    
    # 计算替换区域的尺寸
    w_overlay = int(round(bbox_px[2] - bbox_px[0]))
    h_overlay = int(round(bbox_px[3] - bbox_px[1]))
    
    if w_overlay <= 0 or h_overlay <= 0:
        raise ValueError("边界框尺寸无效")
    
    # 准备要插入的图片
    overlay_bytes = crop_and_resize_to_bytes(image_path, w_overlay, h_overlay)
    
    # 插入图片到指定区域
    page.insert_image(rect, stream=overlay_bytes, overlay=True, keep_proportion=False)
    
    # 保存文档
    safe_overwrite_pdf(doc, output_path)
    doc.close()
    
    print(f"成功替换图片，保存为: {output_path}")

def main():
    # 检查文件是否存在
    if not os.path.exists(PDF_PATH):
        print(f"错误: PDF文件不存在: {PDF_PATH}")
        return
        
    if not os.path.exists(IMAGE_PATH):
        print(f"错误: 图片文件不存在: {IMAGE_PATH}")
        return
        
    if not os.path.exists(JSON_PATH):
        print(f"错误: JSON文件不存在: {JSON_PATH}")
        return
    
    # 加载JSON数据
    try:
        with open(JSON_PATH, 'r', encoding='utf-8') as f:
            json_data = json.load(f)
    except Exception as e:
        print(f"加载JSON文件失败: {e}")
        return
    
    # 在JSON中查找bbox_page_pixel
    try:
        bbox_px, page_index = find_bbox_in_json(json_data, PAPER_ID, PANEL_ID)
        if not bbox_px:
            raise ValueError("在JSON中未找到bbox_page_pixel")
        
        print(f"找到坐标: bbox={bbox_px}, page_index={page_index+1}")
    except Exception as e:
        print(f"查找坐标失败: {e}")
        return
    
    # 执行替换
    try:
        replace_image_in_pdf(
            PDF_PATH, 
            IMAGE_PATH, 
            bbox_px, 
            page_index, 
            OUTPUT_PATH
        )
    except Exception as e:
        print(f"替换失败: {e}")

if __name__ == "__main__":
    main()